import _thread as thread
import base64
import datetime
import hashlib
import hmac
import json
import ssl
from datetime import datetime
from time import mktime
from urllib.parse import urlparse, urlencode
from wsgiref.handlers import format_date_time

import websocket


class ImageUnderstandingClient:
    def __init__(self, appid, apikey, apisecret, image_data, result_callback):
        self.appid = appid
        self.apikey = apikey
        self.apisecret = apisecret
        self.image_data = image_data
        self.imageunderstanding_url = "wss://spark-api.cn-huabei-1.xf-yun.com/v2.1/image"
        self.ws = None
        self.text = [
            {"role": "user", "content": str(base64.b64encode(self.image_data), 'utf-8'), "content_type": "image"},
            {"role": "user",
             "content": "假如你是专业的中医，你简要进行望闻问切中的“望”，记住你是专业的中医，除了分析“望”，不用总结，不要说其他废话",
             "content_type": "text"}
        ]
        self.result_callback = result_callback
        self.message_buffer = []

    class Ws_Param:
        def __init__(self, APPID, APIKey, APISecret, imageunderstanding_url):
            self.APPID = APPID
            self.APIKey = APIKey
            self.APISecret = APISecret
            self.host = urlparse(imageunderstanding_url).netloc
            self.path = urlparse(imageunderstanding_url).path
            self.ImageUnderstanding_url = imageunderstanding_url

        def create_url(self):
            now = datetime.now()
            date = format_date_time(mktime(now.timetuple()))

            signature_origin = f"host: {self.host}\ndate: {date}\nGET {self.path} HTTP/1.1"
            signature_sha = hmac.new(self.APISecret.encode('utf-8'), signature_origin.encode('utf-8'),
                                     digestmod=hashlib.sha256).digest()
            signature_sha_base64 = base64.b64encode(signature_sha).decode('utf-8')

            authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
            authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode('utf-8')

            v = {
                "authorization": authorization,
                "date": date,
                "host": self.host
            }
            url = self.ImageUnderstanding_url + '?' + urlencode(v)
            return url

    def on_error(self, ws, error):
        print(f"### error: {error}")
        self.result_callback(f"Error: {error}")

    def on_close(self, ws, close_status_code, close_msg):
        print(f"### closed ###: {close_status_code}, {close_msg}")
        if self.message_buffer:
            result = ''.join(self.message_buffer).replace("分析结果: ", "").strip()
            self.result_callback(result)

    def on_open(self, ws):
        def run(*args):
            try:
                data = json.dumps(self.gen_params())
                ws.send(data)
            except Exception as e:
                print(f"Error during on_open: {e}")
                self.result_callback(f"Error: {e}")

        thread.start_new_thread(run, ())

    def on_message(self, ws, message):
        data = json.loads(message)
        code = data['header']['code']
        if code != 0:
            print(f'请求错误: {code}, {data}')
            ws.close()
        else:
            choices = data["payload"]["choices"]
            status = choices["status"]
            content = choices["text"][0]["content"]
            print(content, end="")
            self.result_callback(content)
            if status == 2:
                ws.close()

    def gen_params(self):
        data = {
            "header": {
                "app_id": self.appid
            },
            "parameter": {
                "chat": {
                    "domain": "image",
                    "temperature": 0.5,
                    "top_k": 4,
                    "max_tokens": 2028,
                    "auditing": "default",
                    # "streaming": False
                }
            },
            "payload": {
                "message": {
                    "text": self.text
                }
            }
        }
        return data

    def connect(self):
        try:
            ws_param = self.Ws_Param(self.appid, self.apikey, self.apisecret, self.imageunderstanding_url)
            ws_url = ws_param.create_url()
            print(f"WebSocket URL: {ws_url}")
            self.ws = websocket.WebSocketApp(ws_url,
                                             on_message=self.on_message,
                                             on_error=self.on_error,
                                             on_close=self.on_close)
            self.ws.on_open = self.on_open
            self.ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
        except Exception as e:
            print(f"Error during connect: {e}")
            self.result_callback(f"Error: {e}")

